//Christian: Genetic statistics calculator; Simcoal compatable version
#include "genstat.h"

void CalcStats(THapArray *haps, summary_stats *st, const Mut_Type data_type) {
	float a1=0, a2=0, ss=0.0, c2, e1S, e2S, n=0;
	float num_bases=haps[0][0].hap->end()-haps[0][0].hap->begin();
	bool is_seg;
   int maxmis=0,c;
   st->num_haps=haps->GetItemsInContainer();
   st->num_private_haps=0;
   st->mismatch.Flush();
	for(int i=0;i<st->num_haps;i++) {
        n+=haps[0][i].n;
        if(haps[0][i].is_private) st->num_private_haps++;
   }
   st->pairws_dif=0.0;
   //Calculate diversity, ...
   st->hap_div=1.0;
   int mism[st->num_haps-1];
   st->mismatch.Add(0);
	for(int i=0; i<st->num_haps; i++) {
		st->hap_div-=pow((float)(haps[0][i].n/n),2.0);
		for(int j=0; j<st->num_haps-i-1; j++) mism[j]=0; //clear out mism
		//...segregating sites, pairwise dif...
		if(data_type==DNA) {
         for(int curbase=0; curbase<num_bases; curbase++) {
            c=*(haps[0][i].hap->begin()+curbase);
            is_seg=0;
            for(int j=i+1; j<st->num_haps; j++) {
               if(c!= *(haps[0][j].hap->begin()+curbase))
               {  st->pairws_dif+=haps[0][i].n*haps[0][j].n;
                  if(i==0) is_seg=1;
                  mism[j-i-1]++;
               }
            }
            ss+=(float)is_seg;
       	}
       	st->mismatch[0]+=(haps[0][i].n-1)*haps[0][i].n;
       	for(int k=i+1;k<st->num_haps; k++) {
            if(mism[k-i-1]>maxmis) {
               for(int j=maxmis; j<mism[k-i-1]; j++)
                  st->mismatch.Add(0);
               maxmis=mism[k-i-1];
            }
            st->mismatch[mism[k-i-1]]+=haps[0][i].n*haps[0][k].n;
         }
      } else { 
         a1+=haps[0][i].n*pow((float)(*(haps[0][i].hap->begin())),2.0);
         a2+=haps[0][i].n*(float)(*(haps[0][i].hap->begin()));
      }
	}
	if(data_type==DNA) {
      if(n>1.0) {
         st->pairws_dif*=2/(n*(n-1));
         st->nuc_div=st->pairws_dif/num_bases;
         //...and Tajima's D.
       	for(int i=0; i<n; i++)
      		if(i>0) { a1+=1/(float)i; a2+=1/pow((float)i,2.0); }
         e1S=ss*((n+1)/(3*(n-1))-1/a1)/a1;
         c2=2*(pow(n,2.0)+n+3)/(9*n*(n-1))-(n+2)/(a1*n)+a2/pow(a1,2.0);
         e2S=c2*ss*(ss-1)/(pow(a1,2)+a2);
         st->tajd=(st->pairws_dif-ss/a1)/sqrt(e1S+e2S);
         st->seg_sites=(int)ss;
      } else { st->seg_sites=0; st->tajd=0.0; st->nuc_div=0.0;}
   } else { //store allelic variance in pairws_dif for now
      st->pairws_dif=(a1-pow(a2,2.0)/n)/(n-1);
   }
}

void CalcInterStats(THapArray *h1, THapArray *h2, summary_stats *st, const Mut_Type data_type) {
   int ind, n1=0, n2=0;
	float num_bases=h1[0][0].hap->end()-h1[0][0].hap->begin();
	float het1=1.0, het2=1.0;
   st->nuc_div=1.0;
   st->num_haps=0; 
   for(int i=0; i<h1->GetItemsInContainer(); i++) n1+=h1[0][i].n;
   st->num_private_haps=0; 
   for(int i=0; i<h2->GetItemsInContainer(); i++) n2+=h2[0][i].n;
   if(data_type==DNA) {
      for(int i=0; i<h1->GetItemsInContainer(); i++) {
         het1-=pow(float(h1[0][i].n)/float(n1),2.0);
         ind=h2->Find(h1[0][i]);
         if(ind==INT_MAX) { 
            st->num_haps++; //put priv1 in num_haps
            st->nuc_div-=pow(.5*float(h1[0][i].n)/float(n1),2.0);   //pooled heterozygosity (Ht) in nuc_div
         } else
            st->nuc_div-=pow(float(h1[0][i].n)/float(n1)+float(h2[0][ind].n)/float(n2),2)/4.0;
         for(int j=0; j<h2->GetItemsInContainer(); j++) {
            if(i==0) {
               ind=h1->Find(h2[0][j]);
               if(ind==INT_MAX) {
                  st->num_private_haps++; //put priv2 in num_private_haps
                  st->nuc_div-=pow(.5*float(h2[0][j].n)/float(n2),2.0);   //continue on with (Ht)
               }
               het2-=pow(float(h2[0][j].n)/float(n2),2.0);
            }
            for(int curbase=0; curbase<num_bases; curbase++) {
               if(*(h1[0][i].hap->begin()+curbase)!=*(h2[0][j].hap->begin()+curbase))
                  st->pairws_dif+=h1[0][i].n*h2[0][j].n;
            }
         }
      }
      st->pairws_dif*=1.0/(float(n1*n2));
      st->hap_div=(het1+het2)/2;  //put mean expected het (Hs-bar) in hap_div
      st->tajd=(st->nuc_div-st->hap_div)/st->nuc_div;
   }//end if DNA   
}

TPrior* BayesRead(void *buf,istream *ifs,int type) {
//type: 1=int 2=long 3=float
   char c=ifs->peek();
   while(c==' ' || c=='\t') { ifs->get(c); c=ifs->peek(); }
   if(c=='[') {
      TPrior *t=new TPrior;
      (*ifs) >> (*t);
      t->SetLoc1(buf);
      switch(type) {
         case 1: *(int *)buf=int(t->pick()+.5); break;
         case 2: *(long *)buf=long(t->pick()+.5); break;
         case 3: *(float *)buf=(float)t->pick(); break;
      }
   } else {
      switch(type) {
         case 1: (*ifs) >> *(int*)buf; break;
         case 2: (*ifs) >> *(long*)buf; break;
         case 3: (*ifs) >> *(float*)buf; break;
      }
      return NULL;
   }
}

extern istream& operator>>(istream& is, TPrior& P) {
   char buf[50], *x;
   is>>buf;
   x=strtok(buf+1,":");
   P.min=atof(x);
   x=strtok(NULL,"]");
   P.max=atof(x);
   return is;
}

